import cv2
import numpy as np
from PIL import Image, ImageEnhance
from skimage import measure, morphology
from typing import Tuple, Optional
import logging

class ImageProcessor:
    """图像预处理模块"""
    
    def __init__(self, max_size=4000):
        self.logger = logging.getLogger(__name__)
        self.max_size = max_size
    
    def resize_image(self, image: np.ndarray) -> Tuple[np.ndarray, float]:
        """按比例缩放图像，使其最长边不超过max_size"""
        h, w = image.shape[:2]
        if max(h, w) <= self.max_size:
            return image, 1.0

        scale = self.max_size / max(h, w)
        new_w, new_h = int(w * scale), int(h * scale)
        
        self.logger.info(f"图像尺寸过大 ({w}x{h})，正在缩放到 ({new_w}x{new_h})")
        
        resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_AREA)
        return resized, scale
    
    def load_image(self, image_path: str) -> np.ndarray:
        """加载图像文件"""
        try:
            image = cv2.imread(image_path)
            if image is None:
                raise ValueError(f"无法加载图像: {image_path}")
            return image
        except Exception as e:
            self.logger.error(f"加载图像失败: {e}")
            raise
    
    def detect_flowchart_region(self, image: np.ndarray) -> np.ndarray:
        """检测流程图区域"""
        # 转换为灰度图
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # 使用自适应阈值进行二值化，以更好地处理不同光照
        binary = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                                     cv2.THRESH_BINARY_INV, 11, 2)
        
        # 查找轮廓
        contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        if contours:
            # 过滤掉过小的噪点轮廓
            significant_contours = [c for c in contours if cv2.contourArea(c) > 100]
            
            if significant_contours:
                # 计算包含所有重要轮廓的边界框
                all_points = np.vstack([c for c in significant_contours])
                x, y, w, h = cv2.boundingRect(all_points)
                
                # 添加边距
                margin = 50 # 针对大图增加边距
                x = max(0, x - margin)
                y = max(0, y - margin)
                w = min(image.shape[1] - x, w + 2 * margin)
                h = min(image.shape[0] - y, h + 2 * margin)
                
                self.logger.info(f"检测到流程图区域: x={x}, y={y}, w={w}, h={h}")
                return image[y:y+h, x:x+w]
        
        self.logger.warning("未检测到有效的流程图区域，返回原图。")
        return image
    
    def enhance_image_quality(self, image: np.ndarray) -> np.ndarray:
        """增强图像质量"""
        # 转换为PIL图像进行增强
        pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        
        # 对比度增强
        enhancer = ImageEnhance.Contrast(pil_image)
        pil_image = enhancer.enhance(1.5)
        
        # 锐度增强
        enhancer = ImageEnhance.Sharpness(pil_image)
        pil_image = enhancer.enhance(1.2)
        
        # 转换回OpenCV格式
        enhanced = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
        return enhanced
    
    def correct_skew(self, image: np.ndarray) -> Tuple[np.ndarray, float]:
        """自动倾斜校正"""
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # 使用高斯模糊减少噪声，提高边缘检测效果
        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
        
        # 边缘检测
        edges = cv2.Canny(blurred, 50, 150, apertureSize=3)
        
        # 霍夫变换检测直线
        lines = cv2.HoughLinesP(edges, 1, np.pi / 180, threshold=100, minLineLength=100, maxLineGap=10)
        
        if lines is None:
            self.logger.info("未检测到足够长的直线用于倾斜校正。")
            return image, 0.0

        angles = []
        for line in lines:
            x1, y1, x2, y2 = line[0]
            # 避免除以零
            if x2 - x1 == 0:
                continue
            angle = np.degrees(np.arctan((y2 - y1) / (x2 - x1)))
            angles.append(angle)
            
        if not angles:
            self.logger.info("未计算出任何有效角度。")
            return image, 0.0

        # 使用中位数来抵抗异常值
        median_angle = np.median(angles)
        
        # 如果倾斜角度很小，则不进行校正
        if abs(median_angle) < 0.5:
            self.logger.info(f"倾斜角度 {median_angle:.2f}° 过小，无需校正。")
            return image, 0.0
            
        self.logger.info(f"检测到倾斜角度: {median_angle:.2f}°，正在校正...")
        
        # 计算旋转矩阵并执行旋转
        height, width = image.shape[:2]
        center = (width // 2, height // 2)
        rotation_matrix = cv2.getRotationMatrix2D(center, median_angle, 1.0)
        
        corrected = cv2.warpAffine(image, rotation_matrix, (width, height),
                                   flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE)
        
        return corrected, median_angle
    
    def denoise_image(self, image: np.ndarray) -> np.ndarray:
        """图像去噪"""
        # 双边滤波去噪
        denoised = cv2.bilateralFilter(image, 9, 75, 75)
        
        # 形态学操作去除小噪点
        kernel = np.ones((2, 2), np.uint8)
        denoised = cv2.morphologyEx(denoised, cv2.MORPH_CLOSE, kernel)
        
        return denoised
    
    def binarize_image(self, image: np.ndarray) -> np.ndarray:
        """图像二值化"""
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # 使用自适应阈值以获得更好的局部对比度效果
        binary = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                       cv2.THRESH_BINARY_INV, 15, 4)
        
        # 反转图像，使图形为黑色，背景为白色，以匹配后续处理步骤的期望
        # (通常findContours找白色物体，但后续步骤可能期望黑色)
        # 经过验证，后续步骤期望的是黑色背景、白色图形
        
        return binary
    
    def preprocess_image(self, image_path: str) -> Tuple[np.ndarray, np.ndarray, dict]:
        """完整的图像预处理流程"""
        self.logger.info(f"开始预处理图像: {image_path}")
        
        # 加载图像
        image = self.load_image(image_path)
        original_size = image.shape[:2]

        # 智能缩放
        image, scale = self.resize_image(image)

        # 检测流程图区域
        cropped = self.detect_flowchart_region(image)
        
        # 增强图像质量
        enhanced = self.enhance_image_quality(cropped)
        
        # 倾斜校正
        corrected, skew_angle = self.correct_skew(enhanced)
        
        # 去噪
        denoised = self.denoise_image(corrected)
        
        # 二值化
        binary = self.binarize_image(denoised)
        
        metadata = {
            'original_size': original_size,
            'processed_size': corrected.shape[:2],
            'scale_factor': scale, # 记录缩放比例
            'skew_angle': skew_angle,
            'quality_score': self._calculate_quality_score(binary)
        }
        
        self.logger.info(f"图像预处理完成，质量评分: {metadata['quality_score']:.2f}")
        
        return corrected, binary, metadata
    
    def _calculate_quality_score(self, binary_image: np.ndarray) -> float:
        """计算图像质量评分"""
        # 计算连通区域数量
        labeled = measure.label(binary_image)
        regions = measure.regionprops(labeled)
        
        # 基于连通区域数量和分布计算质量评分
        if len(regions) > 0:
            areas = [region.area for region in regions]
            avg_area = np.mean(areas)
            std_area = np.std(areas)
            
            # 质量评分：连通区域数量适中，面积分布均匀
            score = min(1.0, len(regions) / 50) * (1 - std_area / avg_area if avg_area > 0 else 0)
            return max(0.0, score)
        
        return 0.0 